import pandas as pd
from nltk.sentiment.vader import SentimentIntensityAnalyzer
from snownlp import SnowNLP

sia = SentimentIntensityAnalyzer()


def snownlp_sa(sentence):
    s = SnowNLP(sentence)
    return s.sentiments


def calc_df_sa_snownlp():
    covid = 'clean.csv'
    df = pd.read_csv(covid, index_col=0)
    df['sa'] = df['text'].apply(snownlp_sa)
    df.to_csv('sa.csv')
    return df


def nltk_sa_score(sentence):
    sentiment_dict = sia.polarity_scores(sentence)
    return sentiment_dict['pos'] - sentiment_dict['neg']


def nltk_sa_classification(sentence):
    sentiment_dict = sia.polarity_scores(sentence)
    score = sentiment_dict['compound']
    if score >= 0.05:
        return 1
    elif score <= -0.05:
        return -1
    else:
        return 0


def calc_df_sa_nltk():
    covid = 'clean.csv'
    df = pd.read_csv(covid, index_col=0)
    df['sa'] = df['text'].apply(nltk_sa_score)
    df['sa_label'] = df['text'].apply(nltk_sa_classification)
    df.to_csv('sa.csv')
    return df


if __name__ == '__main__':
    df = calc_df_sa_nltk()
